# 简化的脑电数据采集与滤波程序
# 功能：连接NeuroSky设备，采集原始脑电数据，进行滤波处理并记录日志
# 方法：使用NeuroPy3库获取数据，MNE-Python进行信号处理
# 原理：缓存原始数据到指定长度后进行批量滤波处理

import numpy as np
import mne
import time
import csv
import sys
import os
from datetime import datetime

# 添加当前目录到路径以导入NeuroPy3
sys.path.append(os.path.dirname(os.path.abspath(__file__)))
from NeuroPy3 import NeuroPy

class EEGProcessor:
    def __init__(self, port="COM5", buffer_size=512):  # 5秒数据
        self.port = port
        self.buffer_size = buffer_size
        self.data_buffer = []
        self.device = NeuroPy(port, 57600)
        self.log_file = "eeg_analysis_log.csv"
        
        # MNE参数
        self.sfreq = 512
        self.info = mne.create_info(['EEG'], self.sfreq, ['eeg'])
        
        # 初始化日志文件
        self._init_log()
        
    def _init_log(self):
        """初始化CSV日志文件"""
        with open(self.log_file, 'w', newline='', encoding='utf-8') as f:
            writer = csv.writer(f)
            writer.writerow(['时间戳', '数据长度', 'Delta功率', 'Theta功率', 'Alpha功率', 'Beta功率', 'Gamma功率', '信号质量'])
    
    def _data_callback(self, raw_value):
        """数据回调函数"""
        self.data_buffer.append(raw_value)
        
    def _process_data(self):
        """处理缓冲区数据"""
        if len(self.data_buffer) < self.buffer_size:
            return
            
        # 获取数据并清空缓冲区
        data = np.array(self.data_buffer[:self.buffer_size])
        del self.data_buffer[:self.buffer_size]
        
        # 转换单位并创建MNE对象
        data_volts = data / 1e6
        raw = mne.io.RawArray([data_volts], self.info, verbose=False)
        
        # **优化: 过滤和PSD计算**
        # 在计算PSD时，内部可以处理滤波。
        # 让我们计算 1Hz 到 40Hz 的全PSD
        spectrum = raw.compute_psd(method='welch', fmin=1.0, fmax=40.0, verbose=False)
        psds, freqs = spectrum.get_data(return_freqs=True, picks='eeg')
        psds_db = 10 * np.log10(psds) # 转换为分贝单位

        # **新的功率计算方法：在PSD结果中寻找频段**
        delta_power = self._get_band_power_from_psd(psds_db, freqs, 1, 4)
        theta_power = self._get_band_power_from_psd(psds_db, freqs, 4, 8)
        alpha_power = self._get_band_power_from_psd(psds_db, freqs, 8, 13)
        beta_power = self._get_band_power_from_psd(psds_db, freqs, 13, 30)
        gamma_power = self._get_band_power_from_psd(psds_db, freqs, 30, 40)
        
        # **优化: 使用更可靠的信号质量**
        # 假设您已经获取了设备的poorSignal值，将其记录下来
        # quality = self.device.poorSignal
        # 否则，您的原始计算方法作为参考也可以
        quality = self._calculate_quality(data)
        
        # 记录日志
        self._log_result(len(data), delta_power, theta_power, alpha_power, beta_power, gamma_power, quality)
        
        print(f"处理完成 - 数据点: {len(data)}, Delta: {delta_power:.2f} dB, Alpha: {alpha_power:.2f} dB, 质量: {quality:.2f}")

        
    def _get_band_power_from_psd(self, psds, freqs, low_freq, high_freq):
        """从完整的PSD数组中提取指定频段的平均功率（dB单位）"""
        band_idx = np.logical_and(freqs >= low_freq, freqs <= high_freq)
        return psds[0, band_idx].mean()

    def _calculate_band_power(self, raw, low_freq, high_freq):
        """计算指定频段的功率"""
        # 计算功率谱密度
        psd = raw.compute_psd(fmin=low_freq, fmax=high_freq, verbose=False)
        # 返回平均功率
        return np.mean(psd.get_data())
        
    def _calculate_quality(self, data):
        """计算信号质量评分"""
        std_dev = np.std(data)
        return min(100, max(0, 100 - std_dev / 100))
        
    def _log_result(self, length, delta, theta, alpha, beta, gamma, quality):
        """记录分析结果到日志"""
        timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        with open(self.log_file, 'a', newline='', encoding='utf-8') as f:
            writer = csv.writer(f)
            writer.writerow([timestamp, length, f"{delta:.6f}", f"{theta:.6f}", f"{alpha:.6f}", f"{beta:.6f}", f"{gamma:.6f}", f"{quality:.2f}"])
    
    def start(self):
        """启动数据采集"""
        self.device.setCallBack("rawValue", self._data_callback)
        self.device.start()
        print(f"设备已连接到 {self.port}，开始采集数据...")
        
        try:
            while True:
                self._process_data()
                time.sleep(0.1)
        except KeyboardInterrupt:
            print("\n停止采集...")
        finally:
            self.device.stop()
            print("设备已断开")

if __name__ == "__main__":
    processor = EEGProcessor()
    processor.start()